"""
LSTM预测交易策略(Zipline实现)
策略逻辑：
1. 模拟LSTM预测未来价格（实际实现需预训练模型）
2. 结合预测结果生成交易信号
3. 动态调整仓位大小
"""

from zipline.api import order_target_percent, record, symbol, get_datetime
from zipline.finance import commission, slippage
import numpy as np
import pandas as pd

def initialize(context):
    # 策略参数
    context.asset = symbol('AAPL')
    context.lookback = 60  # 输入序列长度
    context.forecast_horizon = 5  # 预测周期
    context.position = 0
    
    # 设置交易成本
    context.set_commission(commission.PerShare(cost=0.001, min_trade_cost=1))
    context.set_slippage(slippage.FixedSlippage(spread=0.01))
    
    # 初始化价格历史
    context.price_history = []

def simulate_lstm_prediction(context, prices):
    """模拟LSTM预测（实际实现需预训练模型）"""
    # 在实际应用中，这里应该使用预训练的LSTM模型进行预测
    # 这里我们使用简单的趋势外推来模拟
    recent_changes = np.diff(prices[-5:])
    avg_change = np.mean(recent_changes) if len(recent_changes) > 0 else 0
    return [prices[-1] + avg_change * (i+1) for i in range(context.forecast_horizon)]

def handle_data(context, data):
    # 获取当前价格
    current_price = data.current(context.asset, 'price')
    context.price_history.append(current_price)
    
    # 保持价格历史长度
    if len(context.price_history) > context.lookback + context.forecast_horizon:
        context.price_history.pop(0)
    
    # 检查是否有足够数据
    if len(context.price_history) < context.lookback + context.forecast_horizon:
        return
    
    # 模拟LSTM预测
    predicted_prices = simulate_lstm_prediction(context, context.price_history)
    avg_predicted_price = np.mean(predicted_prices)
    
    # 计算预期收益率
    expected_return = (avg_predicted_price - current_price) / current_price
    
    # 根据预期收益率确定仓位
    target_position = 0
    if expected_return > 0.05:  # 预期上涨5%以上
        target_position = min(1.0, expected_return * 2)  # 杠杆不超过2倍
    elif expected_return < -0.05:  # 预期下跌5%以上
        target_position = max(-1.0, expected_return * 2)
    
    # 调整仓位
    current_holding = context.portfolio.positions[context.asset].amount
    current_value = current_holding * current_price if current_holding else 0
    current_percent = current_value / context.portfolio.portfolio_value
    
    if abs(target_position - current_percent) > 0.1:  # 避免频繁交易
        order_target_percent(context.asset, target_position)
        context.position = target_position
    
    # 记录状态
    record(price=current_price, 
          predicted_price=avg_predicted_price, 
          expected_return=expected_return, 
          position=context.position)